from django.db import connection, close_old_connections
from django.utils.deprecation import MiddlewareMixin
from django.db.utils import OperationalError
import logging

logger = logging.getLogger(__name__)


class DatabaseConnectionMiddleware(MiddlewareMixin):
    def __init__(self, get_response):
        super().__init__(get_response)
        logger.info("DatabaseConnectionMiddleware initialized...")
        self.get_response = get_response

    def process_request(self, request):
        try:
            connection.cursor().execute("SELECT 1")
        except Exception as e:
            logger.error(f"数据库连接无效，错误信息: {e}，尝试重新连接...")
            try:
                close_old_connections()
                connection.connect()
                logger.info("数据库连接已重新建立")
            except Exception as reconnection_error:
                logger.error(f"重新连接数据库失败，错误信息: {reconnection_error}")
        response = self.get_response(request)
        return response

    def process_response(self, request, response):
        # 在请求处理完成后，这里可以进行一些连接的清理或其他操作
        # 目前示例中暂不做处理，直接返回响应
        return response


class DBCheckMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        try:
            # 尝试执行一个简单的查询来检查连接
            with connection.cursor() as cursor:
                cursor.execute("SELECT 1")
        except OperationalError:
            # 如果连接已经断开，关闭旧连接并重新连接
            close_old_connections()
            connection.connect()
        response = self.get_response(request)
        return response
